"""
市场数据智能体

负责收集和预处理股票市场数据
"""

from agno.agent import Agent
from agno.models.deepseek import DeepSeek
from typing import Dict, Any, Optional
import logging
from datetime import datetime, timedelta
from ..tools.akshare_client import AkShareClient
from ..models.market_data import DataCollectionSummary
from ..models.analysis_results import AgentExecutionLog
import time
logger = logging.getLogger(__name__)


class MarketDataAgent:
    """市场数据智能体类"""

    def __init__(self, model_id: str = "deepseek-chat"):
        """
        初始化市场数据智能体

        Args:
            model_id: LLM模型ID
        """
        self.akshare_client = AkShareClient()

        # 定义数据收集工具
        def collect_stock_data(agent: Agent, ticker: str, start_date: str, end_date: str = "", news_count: int = 5) -> str:
            """
            收集股票综合数据

            Args:
                ticker: 股票代码
                start_date: 开始日期 (YYYYMMDD)
                end_date: 结束日期 (YYYYMMDD)，默认为今天
                news_count: 新闻数量

            Returns:
                str: 数据收集结果摘要
            """
            try:
                if not end_date:
                    end_date = datetime.now().strftime("%Y%m%d")

                logger.info(f"开始收集股票 {ticker} 的数据")

                # 使用AkShare客户端收集数据
                summary = self.akshare_client.collect_comprehensive_data(
                    ticker=ticker,
                    start_date=start_date,
                    end_date=end_date,
                    news_count=news_count
                )

                if "collected_data" not in agent.workflow_session_state:
                    agent.workflow_session_state["collected_data"] = {}

                agent.workflow_session_state["collected_data"][ticker] = {
                    "summary": summary.model_dump(),
                    "collection_time": datetime.now().isoformat(),
                    "status": "completed"
                }

                result_message = f"""数据收集完成！

股票代码: {ticker}
时间范围: {start_date} 至 {end_date}
收集状态: {summary.summary}

详细信息:
- 价格数据: {summary.price_records_count} 条记录
- 新闻数据: {summary.news_count} 条记录
- 数据完整性: {sum(summary.data_collected.values())}/{len(summary.data_collected)} 项成功

数据已保存到智能体状态中，可供后续分析使用。"""

                logger.info(f"股票 {ticker} 数据收集完成")
                return result_message

            except Exception as e:
                error_msg = f"数据收集失败: {str(e)}"
                logger.error(error_msg)

                # 记录错误状态
                if "collected_data" not in agent.workflow_session_state:
                    agent.workflow_session_state["collected_data"] = {}

                agent.workflow_session_state["collected_data"][ticker] = {
                    "status": "failed",
                    "error": str(e),
                    "collection_time": datetime.now().isoformat()
                }

                return error_msg

        
        def get_data_status(agent: Agent) -> str:
            """
            获取数据收集状态
            
            Returns:
                str: 数据收集状态摘要
            """
            if "collected_data" not in agent.workflow_session_state:
                return "尚未收集任何数据"
            
            collected_data = agent.workflow_session_state["collected_data"]
            status_lines = ["数据收集状态摘要:"]
            
            for ticker, data in collected_data.items():
                status = data.get("status", "unknown")
                collection_time = data.get("collection_time", "unknown")
                
                if status == "completed":
                    summary = data.get("summary", {})
                    price_count = summary.get("price_records_count", 0)
                    news_count = summary.get("news_count", 0)
                    status_lines.append(f"- {ticker}: ✅ 完成 (价格:{price_count}条, 新闻:{news_count}条) [{collection_time}]")
                elif status == "failed":
                    error = data.get("error", "unknown error")
                    status_lines.append(f"- {ticker}: ❌ 失败 ({error}) [{collection_time}]")
                else:
                    status_lines.append(f"- {ticker}: ⏳ {status} [{collection_time}]")
            
            return "\n".join(status_lines)
        
        def clear_data_cache(agent: Agent) -> str:
            """
            清除数据缓存
            
            Returns:
                str: 清除结果
            """
            if "collected_data" in agent.workflow_session_state:
                count = len(agent.workflow_session_state["collected_data"])
                agent.workflow_session_state["collected_data"] = {}
                return f"已清除 {count} 个股票的数据缓存"
            else:
                return "数据缓存为空，无需清除"
        
        # 创建Agno智能体
        self.agent = Agent(
            name="Market Data Agent",
            role="股票市场数据收集和预处理专家",
            model=DeepSeek(id=model_id),
            tools=[collect_stock_data, get_data_status, clear_data_cache],
            instructions="""你是一个专业的市场数据收集智能体，负责：

1.  数据收集任务
   - 使用AkShare接口收集A股市场数据
   - 获取股票价格、财务指标、新闻等多维度信息
   - 确保数据的完整性和准确性

2.  数据预处理
   - 验证数据质量和完整性
   - 标准化数据格式
   - 生成数据收集摘要报告

3.  状态管理
   - 维护数据收集状态
   - 提供数据收集进度查询
   - 管理数据缓存和清理

4.  工作流程
   - 接收股票代码和时间范围参数
   - 执行数据收集任务
   - 返回详细的收集结果报告
   - 为后续分析智能体提供数据基础

当前数据收集状态: {collected_data}

请根据用户需求执行相应的数据收集任务，确保数据质量和完整性。""",
            add_state_in_messages=True,
            show_tool_calls=True,
            markdown=True
        )
        
        logger.info("市场数据智能体初始化完成")
    
    def run(self, message: str, **kwargs) -> str:
        """
        运行智能体
        
        Args:
            message: 用户消息
            **kwargs: 其他参数
            
        Returns:
            str: 智能体响应
        """
        try:
            start_time = datetime.now()
            
            # 运行智能体
            response = self.agent.run(message, **kwargs)
            
            end_time = datetime.now()
            duration = (end_time - start_time).total_seconds()
            
            # 记录执行日志
            execution_log = AgentExecutionLog(
                agent_name="MarketDataAgent",
                execution_time=start_time,
                input_data={"message": message, "kwargs": kwargs},
                output_data={"response": response.content if response else "No response"},
                execution_duration=duration,
                status="success" if response else "error",
                error_message=None if response else "No response received"
            )
            
            logger.info(f"市场数据智能体执行完成，耗时: {duration:.2f}秒")
            
            # 正确处理RunResponse对象
            if response:
                if hasattr(response, 'content'):
                    return str(response.content) if response.content else "执行失败，未收到响应内容"
                else:
                    return str(response) if response else "执行失败，响应为空"
            else:
                return "执行失败，未收到响应"
            
        except Exception as e:
            error_msg = f"市场数据智能体执行失败: {str(e)}"
            logger.error(error_msg)
            
            # 记录错误日志
            execution_log = AgentExecutionLog(
                agent_name="MarketDataAgent",
                execution_time=datetime.now(),
                input_data={"message": message, "kwargs": kwargs},
                output_data={},
                execution_duration=0.0,
                status="error",
                error_message=str(e)
            )
            
            return error_msg
    
    def print_response(self, message: str, stream: bool = True, **kwargs):
        """
        打印智能体响应
        
        Args:
            message: 用户消息
            stream: 是否流式输出
            **kwargs: 其他参数
        """
        self.agent.print_response(message, stream=stream, **kwargs)

    def get_analysis_result(self, ticker: str) -> Optional[Dict[str, Any]]:
        """
        获取分析结果
        
        Args:
            ticker: 股票代码
        """
        return self.agent.workflow_session_state.get("collected_data", {}).get(ticker)
